import os
import cv2
import cv2.mat_wrapper
import numpy as np
from tqdm import tqdm
import random

typeName = "val"  # 数据集类型，train 或 val
labels_dir = rf"E:\Python\PythonProject\detect\labels\{typeName}"  # 标签目录
images_dir = rf"E:\Python\PythonProject\detect\dataset\{typeName}\color"  # 图片目录
outputI_dir = rf"E:\Python\PythonProject\GJ\images\{typeName}"  # 保存目录
outputL_dir = rf"E:\Python\PythonProject\GJ\labels\{typeName}"  # 保存目录

os.makedirs(outputI_dir, exist_ok=True)  # 创建输出目录
os.makedirs(outputL_dir, exist_ok=True)  # 创建输出目录

# 创建一个映射表，将原始类别编号映射为新的编号
class_mapping = {2: 0, 3: 1}

# 遍历标签目录中的所有文件
for label_file in tqdm(os.listdir(labels_dir), desc="处理进度", unit="图像"):
    label_path = os.path.join(labels_dir, label_file)
    # 读取标签文件
    with open(label_path, 'r') as f:
        lines = f.readlines()

    # 根据标签文件名找到对应的图片文件
    image_file = label_file.replace('.txt', '.jpg')  # 假设图片是 jpg 格式
    image_path = os.path.join(images_dir, image_file)
    if not os.path.exists(image_path):
        print(f"图片文件 {image_file} 不存在，跳过")
        continue

    # 读取图片
    image = cv2.imread(image_path, 0)
    if image is None:
        print(f"无法读取图片 {image_file}，跳过")
        continue

    # 获取图片的宽度和高度
    height, width = image.shape

    # 创建一个与图片大小相同的黑图
    black_image = np.zeros_like(image)

    # 遍历标签中的每一行
    new_lines = []  # 用于存储修改后的标签行
    for line in lines:
        parts = line.strip().split()
        if len(parts) != 5:
            print(f"标签文件 {label_file} 中的格式不正确，跳过该行")
            continue

        # 提取标注信息
        class_id = int(parts[0])
        x_norm = float(parts[1])
        y_norm = float(parts[2])
        w_norm = float(parts[3])
        h_norm = float(parts[4])

        # 将原始类别编号映射为新的编号
        if class_id in class_mapping:
            new_class_id = class_mapping[class_id]
        else:
            new_class_id = class_id  # 如果不在映射表中，保持原样

        # 将归一化坐标转换为像素坐标
        x = int(x_norm * width)
        y = int(y_norm * height)

        # 在黑图上以标注中心点为中心，30 半径范围内赋值为图片中对应点的像素值
        cv2.circle(black_image, (x, y), 30, 255, -1) 

        # 保存修改后的标签行
        new_line = f"{new_class_id} {x_norm} {y_norm} {w_norm} {h_norm}\n"
        new_lines.append(new_line)
    
    # 随机噪点
    num_noise_points = random.randint(5, 15)  # 随机生成的噪点数量
    for _ in range(num_noise_points):
        # 随机生成 x 和 y 坐标
        random_x = random.randint(30, width-30)
        random_y = random.randint(30, height-30)
        # 在黑图上以随机点为中心，30 半径范围内赋值为 255
        cv2.circle(black_image, (random_x, random_y), 30, 255, -1)

    mask_img = cv2.bitwise_and(image, black_image)
    # 保存修改后的标签文件
    new_label_path = os.path.join(outputL_dir, f"{label_file}")
    with open(new_label_path, 'w') as f:
        f.writelines(new_lines)

    # 保存黑图到输出目录
    output_file = os.path.join(outputI_dir, image_file)
    # cv2.imshow("Mask Image", mask_img)
    # cv2.waitKey(0)
    cv2.imwrite(output_file, mask_img)

    # print(f"黑图已保存到 {output_file}")
    # print(f"修改后的标签文件已保存到 {new_label_path}")